package com.ruoyi.business.alarm.websocket;

import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.WebSocketMessage;
import org.springframework.web.socket.WebSocketSession;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;

/**
 * 报警WebSocket处理器
 * 使用Spring WebSocket API，兼容Undertow
 * 
 * @author ruoyi
 * @date 2024-01-01
 */
@Component
public class AlarmWebSocketHandler implements WebSocketHandler {

    private static final Logger log = LoggerFactory.getLogger(AlarmWebSocketHandler.class);

    /**
     * 存储所有连接的会话
     * key: sessionId
     * value: WebSocketSession
     */
    private static final ConcurrentHashMap<String, WebSocketSession> SESSIONS = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String sessionId = session.getId();
        
        try {
            log.info("收到WebSocket连接请求，sessionId: {}, URI: {}", 
                    sessionId, session.getUri());
            
            // 设置会话超时时间（60秒，与前端心跳超时时间一致）
            // 注意：Spring WebSocket的会话超时时间默认可能较短，需要设置
            // 这里设置一个较长的超时时间，避免因为空闲而断开连接
            // 实际超时由心跳机制控制
            
            // 添加到会话集合
            SESSIONS.put(sessionId, session);
            
            log.info("WebSocket连接建立成功，sessionId: {}, 当前连接数: {}, 所有sessionIds: {}", 
                    sessionId, SESSIONS.size(), SESSIONS.keySet());
            
            // 发送连接成功消息
            String connectedMessage = createMessage("connected", "连接成功", null);
            sendMessage(session, connectedMessage);
            log.info("已发送连接成功消息到sessionId: {}, 消息: {}", sessionId, connectedMessage);
            
        } catch (Exception e) {
            log.error("WebSocket连接建立失败，sessionId: {}", sessionId, e);
            e.printStackTrace();
        }
    }

    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        String sessionId = session.getId();
        
        try {
            if (message instanceof TextMessage) {
                String payload = ((TextMessage) message).getPayload();
                log.debug("收到WebSocket消息，sessionId: {}, message: {}", sessionId, payload);
                
                // 解析消息
                JSONObject jsonMessage = JSON.parseObject(payload);
                String type = jsonMessage.getString("type");
                
                // 处理心跳消息
                if ("ping".equals(type) || "heartbeat".equals(type)) {
                    sendMessage(session, createMessage("pong", "心跳响应", null));
                    return;
                }
                
                // 可以处理其他类型的消息
                // 例如：订阅特定类型的报警、取消订阅等
            }
        } catch (Exception e) {
            log.error("处理WebSocket消息失败，sessionId: {}, message: {}", sessionId, message, e);
        }
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        String sessionId = session.getId();
        log.error("WebSocket发生错误，sessionId: {}, 错误类型: {}, 错误消息: {}", 
                sessionId, exception.getClass().getName(), exception.getMessage(), exception);
        exception.printStackTrace();
        
        try {
            if (session.isOpen()) {
                session.close(CloseStatus.SERVER_ERROR);
            }
        } catch (IOException e) {
            log.error("关闭WebSocket连接失败", e);
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
        String sessionId = session.getId();
        if (sessionId != null) {
            SESSIONS.remove(sessionId);
            log.info("WebSocket连接关闭，sessionId: {}, 关闭状态: {}, 当前连接数: {}", 
                    sessionId, closeStatus, SESSIONS.size());
        }
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }

    /**
     * 发送消息给指定会话
     */
    private void sendMessage(WebSocketSession session, String message) {
        try {
            if (session != null && session.isOpen()) {
                synchronized (session) {
                    session.sendMessage(new TextMessage(message));
                }
            }
        } catch (IOException e) {
            log.error("发送WebSocket消息失败，sessionId: {}", session.getId(), e);
        }
    }

    /**
     * 向所有连接的客户端广播消息
     */
    public static void broadcast(String message) {
        log.info("广播WebSocket消息，连接数: {}, 消息长度: {}", SESSIONS.size(), message.length());
        if (SESSIONS.isEmpty()) {
            log.warn("没有WebSocket连接，无法发送消息");
            return;
        }
        
        AtomicInteger successCount = new AtomicInteger(0);
        AtomicInteger failCount = new AtomicInteger(0);
        
        SESSIONS.forEach((sessionId, session) -> {
            try {
                if (session.isOpen()) {
                    synchronized (session) {
                        session.sendMessage(new TextMessage(message));
                        successCount.incrementAndGet();
                        log.debug("成功发送消息到sessionId: {}", sessionId);
                    }
                } else {
                    // 移除已关闭的会话
                    SESSIONS.remove(sessionId);
                    failCount.incrementAndGet();
                    log.warn("会话已关闭，移除sessionId: {}", sessionId);
                }
            } catch (Exception e) {
                log.error("广播WebSocket消息失败，sessionId: {}", sessionId, e);
                // 移除异常的会话
                SESSIONS.remove(sessionId);
                failCount.incrementAndGet();
            }
        });
        
        log.info("WebSocket消息广播完成，成功: {}, 失败: {}", successCount.get(), failCount.get());
    }

    /**
     * 向所有连接的客户端发送报警消息
     */
    public static void sendAlarm(Object alarmData) {
        log.info("准备发送报警消息，当前连接数: {}", SESSIONS.size());
        String message = createMessage("alarm", "报警通知", alarmData);
        log.debug("报警消息内容: {}", message);
        broadcast(message);
    }

    /**
     * 向所有连接的客户端发送报警结束消息
     */
    public static void sendAlarmEnd(Object alarmData) {
        String message = createMessage("alarm_end", "报警结束", alarmData);
        broadcast(message);
    }

    /**
     * 创建消息JSON字符串
     */
    private static String createMessage(String type, String message, Object data) {
        JSONObject json = new JSONObject();
        json.put("type", type);
        json.put("message", message);
        json.put("data", data);
        json.put("timestamp", System.currentTimeMillis());
        return json.toJSONString();
    }

    /**
     * 获取当前连接数
     */
    public static int getConnectionCount() {
        return SESSIONS.size();
    }
}

